#encoding: utf-8
from logging import log
import multiprocessing
import platform
import argparse
from gensim.models import Word2Vec
from gensim.models.word2vec import LineSentence
from pycw2vec.config.word2vec_config import configs as config
from pycw2vec.utils.logginger import init_logger
from pycw2vec.utils.utils import file_exists
from pyw2c.DataProcess import WiKiPreprocessor


def dealWiki(logger):
    '''处理维基语料'''
    if platform.system().lower() == 'windows':
        config['wiki_ori_data_path'] = "F:/lrr/dataset/zhwiki-latest-pages-articles.xml.bz2"
    wiki_press = WiKiPreprocessor(config=config, logger=logger)
    wiki_press.xml2txt()
    wiki_press.tradition2simple()
    wiki_press.separate()
    wiki_press.remove()


def main():
    logger = init_logger(log_name='word2vec', log_dir=config['log_dir2'])

    if not file_exists(config['wiki_res_data_path']):
        logger.info('开始读取维基原始数据')
        dealWiki(logger=logger)

    logger.info("开始训练word2vec")
    model = Word2Vec(
        sentences=LineSentence(open(config['sougo_res_data_path'])),
        size=config['embedding_dim'],
        window=config['window_size'],
        min_count=config['min_freq'],
        negative=config['negative_sample_num'],
        seed=args['seed'],
        sg=1,
        iter=15,
        workers=multiprocessing.cpu_count()-1
    )
    logger.info("训练完成，开始保存")

    model.save(config['sougo_w2c_model_path'])

    logger.info("保存完成，执行结束")

if __name__ == '__main__':
    ap = argparse.ArgumentParser(description='gensim w2c model training')
    ap.add_argument('-s',
                    '--seed',
                    default=2018,
                    type=str,
                    help='Seed for initializing training.')
    args = vars(ap.parse_args())
    main()
